
use std::{env::current_dir, fs::{create_dir_all, remove_file, File, OpenOptions}, io::{Read, Write}, path::PathBuf};

use clap::Args;
use serde::{Deserialize, Serialize};

use wtask_base::{cmd::file_zip, error::{WError, WResult}, export::Local, tcp::{WTcpStream, stream_read_data, stream_read_raw, stream_write_data, stream_write_raw}};
use crate::{basic::WRaw, file::utils::FileUDArgs, utils::{AppTrait, StreamStatus}};


#[derive(Debug, Clone, Serialize, Deserialize, Args)]
pub struct AppFileDownload {
    /// 是否返回文件原始数据，默认返回本地下载的文件路径，即下载的数据写入本地文件，返回文件数据时不写入本地文件
    #[arg(long)]
    pub raw: bool,

    /// 如果是文件夹，压缩进行下载，默认不压缩文件夹并终止传输
    #[arg(long)]
    pub zip_dir: bool,

    #[command(flatten)]
    pub data: FileUDArgs,

    /// 下载到本地的路径（默认为当前路径）
    pub dir: Option<PathBuf>,
}


impl From<FileUDArgs> for AppFileDownload {
    fn from(value: FileUDArgs) -> Self {
        Self { data: value, raw: false, dir: None, zip_dir: false }
    }
}


impl AppTrait for AppFileDownload {
    type Status = ();
    type Res = WRaw;
    
    fn client_handle_before<M: FnMut(bool, String) + Send + Sync>(mut self, mut msg_func: M, _status: Option<Self::Status>) -> WResult<(Option<Self>, Option<Self::Status>)> {
        if !self.data.file.is_absolute() {
            msg_func(true, format!("File need absolute, stop: {:?}", self.data.file));
            return Ok((None, None));
        }
        let r = if self.raw {
            Some(self)
        } else {
            let dir_download = if let Some(dir) = self.dir {
                if dir.is_absolute() {
                    dir
                } else {
                    current_dir()?.join(dir)
                }
            } else {
                current_dir()?
            };
            let file_download = dir_download.join(self.data.file.file_name().unwrap());
            if file_download.exists() && (!self.data.overwrite) {
                msg_func(true, format!("File exist, stop: {:?}", file_download));
                None
            } else {
                self.dir = Some(dir_download);
                Some(self)
            }
        };
        Ok((r, None))
    }
    
    async fn client_handle<M: FnMut(bool, String) + Send + Sync>(self, stream: &WTcpStream, key: &str, mut msg_func: M, _status: Option<Self::Status>) -> WResult<Self::Res> {
        let file_name = {
            let msg: String = stream_read_data(stream.0.as_ref(), key).await?;
            if msg.starts_with("File") {
                msg_func(true, msg);
                return Err(WError::Stop);
            }
            msg_func(false, format!("Downloading {} ...", msg));
            msg
        };

        let time_start = Local::now();
        let res = if self.raw {
            WRaw::Raw(file_name)
        } else {
            let mut file_size = 0;
            let file_download = if let Some(dir) = self.dir {
                create_dir_all(&dir)?;
                dir
            } else {
                current_dir()?
            }.join(file_name);
            msg_func(false, format!("Saved: {:?}", file_download));
            let mut f = OpenOptions::new()
                .write(true)
                .create(true)
                .truncate(true)
                .open(&file_download)?;
            while let Ok(chunk_data) = stream_read_raw(stream.0.as_ref(), key).await {
                f.write_all(&chunk_data)?;
                file_size += chunk_data.len();
                msg_func(false, format!("Downloading {}mb ...", file_size / 1024 / 1024));
            }
            let seconds_cost = Local::now().signed_duration_since(time_start).as_seconds_f64();
            let speed = file_size as f64 / seconds_cost / 1024.0 / 1024.0;
            msg_func(false, format!("Download complete, {:.1}mb/s ...", speed));
            WRaw::from(&file_download)
        };
        Ok(res)
    }


    async fn server_handle(self, _server_token: &str, session_token: &str, stream_session: &WTcpStream, _stream_status: StreamStatus) -> WResult<()> {
        let (file_download, need_remove) = if self.zip_dir && self.data.file.is_dir() {
            let file_download = PathBuf::from(format!("{}.zip", self.data.file.display()));
            if let Err(e) = file_zip(&self.data.file, &file_download) {
                let msg = format!("File zip error: {}", e);
                stream_write_data(&msg, stream_session.1.as_ref(), session_token).await?;
                return Ok(());
            }
            (file_download, true)
        } else {
            (self.data.file, false)
        };

        let (msg, f) = match File::open(&file_download) {
            Ok(f) => (file_download.file_name().unwrap().to_string_lossy().to_string(), Some(f)),
            Err(e) => (format!("File error: {}", e), None),
        };
        stream_write_data(&msg, stream_session.1.as_ref(), session_token).await?;
        if f.is_none() {
            return Ok(());
        }
        let mut f = f.unwrap();
        let mut chunk_data = vec![0; self.data.chunk_size];
        while let Ok(n) = f.read(&mut chunk_data) {
            if n == 0 {
                break;
            }
            stream_write_raw(stream_session.1.as_ref(), &chunk_data[..n], session_token).await?;
        }
        if need_remove {
            let _ = remove_file(file_download);
        }
        Ok(())
    }
}
